#include "cvxEMTreeSolver.h"
#include "nvmath.h"
#include <conio.h>
#include <time.h>

using cvx::Rgb8UImage;
using cvx::Bw8UImage;
using cvx::Bw32FImage;
using cvx::Bw32SImage;
using cvx::Rgb32FImage;
using nv::vec3f;

//local window size
static const int IS_WSIZE = 3;    //size of the local windows
static const int IS_WSIZE2 = 9;   //WSIZE2 = WSIZE*WSIZE, number of pixels in each window ('n' in the paper)
static const int IS_MSIZE = 12;   //WSIZE2+3, number of rows in the M matrix (see equation 8)
static const int IS_WSIZE_2 = 1;  //WSIZE_2 = WSIZE/2

/************************************************************************/
/*
	CvxEMtreeSolver
*/
/************************************************************************/

CvxEMtreeSolver::CvxEMtreeSolver()
{
	m_ETree = NULL;
	m_MTree = NULL;

}
CvxEMtreeSolver::~CvxEMtreeSolver()
{
	
}

bool CvxEMtreeSolver::SetETree(const Mat &initialShadingMap, const int threshold)
{
	assert(!m_orgImg.empty());
	assert(!m_edgeMap.empty());
	assert(initialShadingMap.depth() == CV_32F);

	int h = m_orgImg.rows;
	int w = m_orgImg.cols;	

	//init known map
	Mat knownMap = cv::Mat::zeros(h, w, CV_8U);
	Bw32FImage init_map(&initialShadingMap);
	Bw8UImage known_map(&knownMap);
	for (int y = 0; y<h; ++y)
	{
		for (int x = 0; x<w; ++x)
		{
			if (init_map[y][x] > 0.001)
			{
				known_map[y][x] = 255;
			}
		}
	}
	
	unsigned int level = CvxQuadTree::GetLevel(std::max(w, h));

	//set E tree
	m_ETree = new CvxAdaptiveQuadTree(level, m_orgImg, threshold);
	m_ETree->create(NULL);
	m_ETree->balance();
	m_ETree->initLeafData(m_orgImg, initialShadingMap, knownMap, NULL, m_edgeMap);

	return true;
}

bool CvxEMtreeSolver::SetETree(const int stddevThreshold, const int sizeThrehold, const float minimumShadingScale)
{
	assert(!m_orgImg.empty());
	assert(!m_edgeMap.empty());

	int h = m_orgImg.rows;
	int w = m_orgImg.cols;	

	unsigned int level = CvxQuadTree::GetLevel(std::max(w, h));

	//set E tree
	m_ETree = new CvxAdaptiveQuadTree(level, m_orgImg, stddevThreshold);
	m_ETree->create(NULL);
	m_ETree->balance();
	m_ETree->adaptiveInitLeafData(m_orgImg, m_edgeMap, NULL, sizeThrehold, minimumShadingScale);

	int leafNum = 0;
	m_ETree->leafNumber(NULL, leafNum); 
	printf("leafNum = %d\n", leafNum);
	return true;
}

bool CvxEMtreeSolver::SolveETree(bool useEdge)
{
	if (m_ETree)
	{
		vector<map<int, double>> leftMap;
		vector<double> rightVec;
		m_ETree->planeEquation(NULL, leftMap, rightVec);
		if(useEdge)
		{
			m_ETree->edgeEquation(NULL, leftMap, rightVec);
		}

		assert(leftMap.size() == rightVec.size());

		int unknownNum = m_ETree->unknownNodeNumber();
		vector<double> result(unknownNum);
		if (unknownNum > 0)
		{	
		//	double tt = clock();
			bool isOk = LeastSquare::solver(leftMap, rightVec, true, unknownNum, &result[0]);
		//	cout<<"E solver cost time = "<<clock()-tt<<endl;
			if (isOk)
			{
				m_ETree->setShadingResult(m_shading, NULL, result);
				return true;
			}
			else
			{
				return false;
			}
		}
		else
		{
			cout<<"unknown number is zero"<<endl;
		}	
		return true;
	}
	else
	{		
		_cprintf("error: set E tree first.\n");
		return false;
	}
}

bool CvxEMtreeSolver::CalcMTree(bool reverse)
{
	assert(m_ETree);
	if (m_MTree == NULL)
	{
		m_MTree = new CvxAdaptiveQuadTree(*m_ETree);
		if (reverse)
		{
			m_MTree->reverse(NULL);
		}
	}
	return true;
}

bool CvxEMtreeSolver::SolveMtree(bool useEdge)
{
	if (m_MTree)
	{
		int unknownNum = m_MTree->unknownNodeNumber();		
		if (unknownNum > 0)
		{
			vector<double> result(unknownNum);
			vector<map<int, double>> leftMap;
			vector<double> rightVec;
			m_MTree->planeEquation(NULL, leftMap, rightVec, true);

			if(useEdge)
			{
				m_MTree->edgeEquation(NULL, leftMap, rightVec);
			}

			assert(leftMap.size() == rightVec.size());

		//	double tt = clock();
		//	cout<<"start M solver"<<endl;
			bool isOk = LeastSquare::solver(leftMap, rightVec, true, unknownNum, &result[0]);
		//	cout<<"M solver cost time = "<<clock()-tt<<endl;
			if (isOk)
			{
				m_MTree->setShadingResult(m_shading, NULL, result);
				return true;
			}
			else
			{
				return false;
			}		
		}
	}
	return false;	
}

bool CvxEMtreeSolver::UpsampleAvector()
{
	assert(m_MTree);

	int h = m_orgImg.rows;
	int w = m_orgImg.cols;	
	int unknownNum = m_MTree->m_nodeNum;

	//calculate A vector
	vector<map<int, double>> leftMap;
	vector<double> rightVec;

	m_MTree->AvectorEquation(NULL, leftMap, rightVec);	

	vector<double> result(unknownNum*3);
	bool isOk = LeastSquare::solver(leftMap, rightVec, true, unknownNum*3, &result[0]);
	if (isOk)
	{
		Mat AvectorMatrix = Mat(h, w, CV_32FC3, Scalar(0, 0, -255, 0));
		m_MTree->setAvector(NULL, result);
		m_MTree->InterpolateAvector(NULL, AvectorMatrix);

		//save previous result
		savePrevious();

		Bw32FImage sha_map(&m_shading);
		Rgb32FImage avector_map(&AvectorMatrix);
		Rgb8UImage org_img(&m_orgImg);
		Rgb8UImage albedo_img(&m_albedo);
		Bw32FImage mini_map(&m_miniShading);

		for (int y = 0; y<h; ++y)
		{
			for (int x =0; x<w; ++x)
			{
				sha_map[y][x] = org_img[y][x].r * avector_map[y][x].r + 
						org_img[y][x].g * avector_map[y][x].g + 
						org_img[y][x].b * avector_map[y][x].b;
				sha_map[y][x] = std::min(0.99f, sha_map[y][x]);
				if (sha_map[y][x] < mini_map[y][x])
				{
					sha_map[y][x] = mini_map[y][x];
				}
				albedo_img[y][x].r = org_img[y][x].r/sha_map[y][x];
				albedo_img[y][x].g = org_img[y][x].g/sha_map[y][x];
				albedo_img[y][x].b = org_img[y][x].b/sha_map[y][x];			
			}
		}

// 		Mat showShading = m_shading.clone();				
// 		showShading.convertTo(showShading, CV_8U, 255.0);
// 		imshow("upsample quad shading", showShading);
// 		imshow("albedo", m_albedo);

		return true;
	}
	return false;
}

bool CvxEMtreeSolver::UpdataETreeEdgeNode()
{
	assert(m_ETree);
	m_ETree->updateEdgeNodeShading(NULL, m_shading, m_edgeMap, m_albedo, 5);
	return true;
}

void CvxEMtreeSolver::SaveQuadTree(const string &fileName, bool isShow)
{
	if(m_ETree)
	{
		m_ETree->draw(NULL);
		imwrite(fileName, m_ETree->m_orgBackUp);

		if(isShow)
		{
			cv::imshow("quad tree", m_ETree->m_orgBackUp);
		}
	}
}
void CvxEMtreeSolver::SaveColorQuadTree(const string &fileName, bool isShow)
{
	if(m_ETree)
	{
		m_ETree->draw(NULL);
		m_ETree->drawColorNode(NULL, Scalar(255, 0, 0), Scalar(0, 0, 255)); 
		imwrite(fileName, m_ETree->m_orgBackUp);

		if(isShow)
		{
			cv::imshow("color quad tree", m_ETree->m_orgBackUp);
		}
	}

}
void CvxEMtreeSolver::SaveShadingImage(const string &fileName, bool isShow)
{
	Mat showShading = m_shading.clone();				
	showShading.convertTo(showShading, CV_8U, 255.0);
	imwrite(fileName, showShading);

	if(isShow)
	{
		cv::imshow("shading image", showShading);
	}
}

void CvxEMtreeSolver::SaveAlbedoImage(const string &fileName, bool isShow)
{
	imwrite(fileName, m_albedo);

	if(isShow)
	{
		cv::imshow("albdeo image", m_albedo);
	}
}

bool CvxEMtreeSolver::SetGroundTruth2EdgeNode(Mat &shadingGroundTruth)
{
	assert(m_ETree);

	m_ETree->setShadingToEdgeNode(NULL, m_edgeMap, shadingGroundTruth);

	return true;

}
